/* 
 * File:   OCRToolkit.cpp
 * Author: liuyi
 * 
 * Created on July 21, 2009, 4:58 PM
 */

#include "OCR.h"
#include "Log.h"
#include "FontLib.h"
#include "ml/KNNTraversal.h"
#include "Debug.h"
#include "ml/SampleSpace.h"
#include "ml/KMeans.h"
#include "ml/MPLSH.h"
#include "Detector.h"
#include "IPToolkit.h"
#include "CharFeature.h"
#include "Config.h"

#include <cstring>
#include <cfloat>
#include <cmath>
#include <list>

KNN * OCR::s_knn = NULL;
SampleSpace * OCR::s_sampleSpace = NULL;

void OCR::train() {
    Config &config = Config::getInstance();

    string knn = config.getString("knn");

    SampleSpace ss(CharFeature::s_FEATURE_SIZE);
    if (knn == "traversal") {
        s_knn = new KNNTraversal(ss);
    } else if (knn == "mplsh") {
        s_knn = new MPLSH(ss);
    } else {
        assert(false);
    }

    vector<string> pathList;
    string path = config.getDataDirPath();;
    string featurePath = path + "/max/heiti_64_feature";
    pathList.push_back(featurePath);
    trainSubset(pathList);
}

void OCR::trainSubset(vector<string> &pathList) {
    string currLibPath = Config::getInstance().getDataDirPath() + "/curr/library";
    
    wifstream wfin(currLibPath.c_str());
    wfin.imbue(locale(Config::getInstance().getLocale()));
    wstring currLib;
    wfin >> currLib;
    wfin.close();

    int size = pathList.size();
    vector<Sample *> *lists = new vector<Sample *>[size];
    SampleSpace ss(CharFeature::s_FEATURE_SIZE);
    for(int i = 0; i < size; i++) {
        ifstream fin(pathList.at(i).c_str());
        ss.loadSampleFromStream(fin, lists[i]);
        fin.close();
    }

    int offset = 0;
    int len = lists[0].size();
    int distorteNum = Config::DISTORTE_NUM;
    vector<Sample *> sampleList;
    for (int i = 0; i < len; i += distorteNum) {
        ushort a = (ushort)currLib[offset];
        ushort b = (ushort)lists[0].at(i)->getClass();

        if(a == b) {
            for(int j = 0; j < size; j++) {
                for(int k = 0; k < distorteNum; k++) {
                    sampleList.push_back(lists[j].at(i + k));
                }
            }

            offset++;
        }
    }

    s_knn->train(sampleList);

    string path = Config::getInstance().getModelPath();
    s_knn->store2File(path.c_str());

    for(int i = 0; i < size; i++) {
        ss.deleteSampleArray(lists[i]);
    }
    delete []lists;
}

void OCR::preRecognize() {
    Config &config = Config::getInstance();

    try {
        string knn = config.getString("knn");

        if (knn == "traversal") {
            s_knn = KNNTraversal::loadFromFile(config.getModelPath().c_str());
        } else if (knn == "mplsh") {
            s_knn = MPLSH::loadFromFile(config.getModelPath().c_str());
        } else {
            assert(false);
        }
    } catch (FileNotFoundException e) {
        Log::log(e.what());
        exit(-1);
    }

    s_knn->preClassify(s_K);

    s_sampleSpace = &s_knn->getSampleSpace();
}

void OCR::release() {
    if(s_knn != NULL) {
        delete s_knn;
        s_knn = NULL;
    }
}

ushort OCR::recognize(const uchar *greyData, int width, int height) {
    if (!s_knn) {
        Log::log("Classifier has not been initialized!", Log::ERROR);
        exit(1);
    }

    int normSize = CharFeature::s_NORM_SIZE;

    uchar* normChar = new uchar[normSize * normSize];
    normalize(normChar, greyData, width, 0, 0, width, height);

#ifdef SHOW_NORM_CHAR
    Debug::showImageData(normChar, normSize, normSize);
#endif

    Sample *sample = s_sampleSpace->newSample();

    CharFeature::getInstance().extractScaledFeature(sample->getFeature(), normChar);

    s_knn->classify(*sample);

    int label = sample->getClass();
    ushort ret = (label == -1) ? (ushort) L'?' : (ushort) label;

    s_sampleSpace->deleteSample(&sample);
    delete[] normChar;

    return ret;
}

void OCR::normalize(uchar* dst, const uchar* src, int sw, int rx, int ry, int rw, int rh) {
    int normSize = CharFeature::s_NORM_SIZE;

    uchar *region = new uchar[rw * rh];
    for (int i = 0; i < rh; i++) {
        memcpy(region + rw*i, src + sw * (i + ry) + rx, rw);
    }

    memset(dst, BG, normSize * normSize * sizeof (uchar));

    float multiple;
    if (rw >= rh) {
        multiple = normSize * 1.0f / rw;
        int tempH = floor(rh * multiple);
        int count = (normSize - tempH) / 2;

        IPToolkit::resize(dst + normSize*count, normSize, tempH, region, rw, rh);
        IPToolkit::binarize(dst + normSize*count, normSize, tempH, 128);
    } else {
        multiple = normSize * 1.0f / rh;
        int tempW = floor(rw * multiple);
        uchar *tempDst = new uchar[tempW * normSize];

        IPToolkit::resize(tempDst, tempW, normSize, region, rw, rh);
        IPToolkit::binarize(tempDst, tempW, normSize, 128);

        int count = (normSize - tempW) / 2;
        for (int i = 0; i < normSize; i++) {
            memcpy(dst + normSize * i + count, tempDst + tempW*i, tempW);
        }

        delete []tempDst;
    }

    delete []region;
}